package com.wfg.ai.core;

import embedding.Embedding;
import embedding.EmbeddingVector;
import splitter.CodeChunk;
import splitter.Splitter;
import splitter.SplitterConfig;
import splitter.SplitterType;
import splitter.AstCodeSplitter;
import splitter.LangChainCodeSplitter;
import vectordb.VectorDB;
import vectordb.VectorDocument;

import java.io.File;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.Collectors;

/**
 * 代码索引器
 * 负责将代码分割成块，生成嵌入向量，并将它们存储到向量数据库中
 */
public class Indexer {
    private Embedding embedding;
    private VectorDB vectorDB;
    private String collectionName;
    private SplitterConfig splitterConfig;
    private Map<String, Splitter> splitters;
    private boolean initialized;

    /**
     * 索引器配置类
     */
    public static class IndexerConfig {
        private Embedding embedding;
        private VectorDB vectorDB;
        private String collectionName;
        private SplitterConfig splitterConfig;

        public IndexerConfig(Embedding embedding, VectorDB vectorDB, String collectionName) {
            this.embedding = embedding;
            this.vectorDB = vectorDB;
            this.collectionName = collectionName;
            this.splitterConfig = new SplitterConfig();
        }

        public IndexerConfig(Embedding embedding, VectorDB vectorDB, String collectionName, SplitterConfig splitterConfig) {
            this.embedding = embedding;
            this.vectorDB = vectorDB;
            this.collectionName = collectionName;
            this.splitterConfig = splitterConfig != null ? splitterConfig : new SplitterConfig();
        }

        public Embedding getEmbedding() {
            return embedding;
        }

        public void setEmbedding(Embedding embedding) {
            this.embedding = embedding;
        }

        public VectorDB getVectorDB() {
            return vectorDB;
        }

        public void setVectorDB(VectorDB vectorDB) {
            this.vectorDB = vectorDB;
        }

        public String getCollectionName() {
            return collectionName;
        }

        public void setCollectionName(String collectionName) {
            this.collectionName = collectionName;
        }

        public SplitterConfig getSplitterConfig() {
            return splitterConfig;
        }

        public void setSplitterConfig(SplitterConfig splitterConfig) {
            this.splitterConfig = splitterConfig;
        }
    }

    /**
     * 创建一个新的代码索引器
     * 
     * @param config 索引器配置
     */
    public Indexer(IndexerConfig config) {
        this.embedding = config.getEmbedding();
        this.vectorDB = config.getVectorDB();
        this.collectionName = config.getCollectionName();
        this.splitterConfig = config.getSplitterConfig();
        this.splitters = new ConcurrentHashMap<>();
        this.initialized = false;
    }

    /**
     * 初始化索引器
     * 
     * @return 完成后的 Future
     */
    public CompletableFuture<Void> init() {
        if (this.initialized) {
            return CompletableFuture.completedFuture(null);
        }

        return this.vectorDB.init()
                .thenCompose(result -> this.vectorDB.hasCollection(this.collectionName))
                .thenCompose(hasCollection -> {
                    if (!hasCollection) {
                        Map<String, String> metadataFields = new HashMap<>();
                        metadataFields.put("relativePath", "VarChar");
                        metadataFields.put("startLine", "Int64");
                        metadataFields.put("endLine", "Int64");
                        metadataFields.put("fileExtension", "VarChar");
                        
                        return this.vectorDB.createCollection(this.collectionName, this.embedding.getDimension(), metadataFields);
                    }
                    return CompletableFuture.completedFuture(null);
                })
                .thenApply(result -> {
                    this.initialized = true;
                    System.out.println("Indexer initialized with collection: " + this.collectionName);
                    return null;
                })
                .exceptionally(e -> {
                    System.err.println("Failed to initialize indexer: " + e.getMessage());
                    throw new RuntimeException("Failed to initialize indexer", e);
                });
    }

    /**
     * 获取适合特定语言的分割器
     * 
     * @param language 语言
     * @return 分割器
     */
    private Splitter getSplitter(String language) {
        if (this.splitters.containsKey(language)) {
            return this.splitters.get(language);
        }

        Splitter splitter;
        if (this.splitterConfig.getType() == SplitterType.AST && AstCodeSplitter.isLanguageSupported(language)) {
            splitter = new AstCodeSplitter(this.splitterConfig.getChunkSize(), this.splitterConfig.getChunkOverlap());
        } else {
            splitter = new LangChainCodeSplitter(this.splitterConfig.getChunkSize(), this.splitterConfig.getChunkOverlap());
        }

        this.splitters.put(language, splitter);
        return splitter;
    }

    /**
     * 从文件路径获取语言
     * 
     * @param filePath 文件路径
     * @return 语言
     */
    private String getLanguageFromPath(String filePath) {
        String extension = "";
        int i = filePath.lastIndexOf('.');
        if (i > 0) {
            extension = filePath.substring(i + 1);
        }

        switch (extension.toLowerCase()) {
            case "js":
                return "javascript";
            case "ts":
            case "tsx":
                return "typescript";
            case "py":
                return "python";
            case "java":
                return "java";
            case "c":
                return "c";
            case "cpp":
            case "cc":
            case "cxx":
            case "h":
            case "hpp":
                return "cpp";
            case "go":
                return "go";
            case "rs":
                return "rust";
            case "php":
                return "php";
            case "rb":
                return "ruby";
            case "swift":
                return "swift";
            case "scala":
                return "scala";
            case "html":
            case "htm":
                return "html";
            case "md":
                return "markdown";
            case "tex":
                return "latex";
            case "sol":
                return "solidity";
            default:
                return extension;
        }
    }

    /**
     * 索引单个文件
     * 
     * @param filePath 文件路径
     * @param relativePath 相对路径
     * @return 完成后的 Future
     */
    public CompletableFuture<Void> indexFile(String filePath, String relativePath) {
        return this.init()
                .thenCompose(result -> {
                    try {
                        String content = Files.readString(Paths.get(filePath));
                        String language = getLanguageFromPath(filePath);
                        String fileExtension = "";
                        int i = filePath.lastIndexOf('.');
                        if (i > 0) {
                            fileExtension = filePath.substring(i + 1);
                        }

                        Splitter splitter = getSplitter(language);
                        
                        return splitter.split(content, language, relativePath)
                                .thenCompose(chunks -> {
                                    List<CompletableFuture<VectorDocument>> documentFutures = new ArrayList<>();
                                    
                                    for (CodeChunk chunk : chunks) {
                                        CompletableFuture<VectorDocument> docFuture = this.embedding.embed(chunk.getContent())
                                                .thenApply(embeddingVector -> {
                                                    String id = UUID.randomUUID().toString();
                                                    return new VectorDocument(
                                                        id,
                                                        embeddingVector.getVector(),
                                                        chunk.getContent(),
                                                        relativePath,
                                                        chunk.getMetadata().getStartLine(),
                                                        chunk.getMetadata().getEndLine(),
                                                        fileExtension
                                                    );
                                                });
                                        
                                        documentFutures.add(docFuture);
                                    }
                                    
                                    return CompletableFuture.allOf(documentFutures.toArray(new CompletableFuture[0]))
                                            .thenApply(v -> documentFutures.stream()
                                                    .map(CompletableFuture::join)
                                                    .collect(Collectors.toList()));
                                })
                                .thenCompose(documents -> {
                                    if (documents.isEmpty()) {
                                        return CompletableFuture.completedFuture(null);
                                    }
                                    
                                    return this.vectorDB.insert(this.collectionName, documents);
                                });
                    } catch (Exception e) {
                        System.err.println("Failed to index file " + filePath + ": " + e.getMessage());
                        return CompletableFuture.completedFuture(null);
                    }
                });
    }

    /**
     * 索引目录
     * 
     * @param dirPath 目录路径
     * @param baseDir 基础目录（用于计算相对路径）
     * @param includePatterns 包含的文件模式
     * @param excludePatterns 排除的文件模式
     * @return 完成后的 Future
     */
    public CompletableFuture<Void> indexDirectory(String dirPath, String baseDir, List<String> includePatterns, List<String> excludePatterns) {
        return this.init()
                .thenCompose(result -> {
                    try {
                        List<String> filePaths = findFiles(dirPath, includePatterns, excludePatterns);
                        AtomicInteger processedFiles = new AtomicInteger(0);
                        int totalFiles = filePaths.size();
                        
                        List<CompletableFuture<Void>> indexFutures = new ArrayList<>();
                        
                        for (String filePath : filePaths) {
                            String relativePath = Paths.get(baseDir).relativize(Paths.get(filePath)).toString();
                            
                            CompletableFuture<Void> indexFuture = indexFile(filePath, relativePath)
                                    .thenApply(v -> {
                                        int processed = processedFiles.incrementAndGet();
                                        System.out.println("Indexed file " + processed + "/" + totalFiles + ": " + relativePath);
                                        return null;
                                    });
                            
                            indexFutures.add(indexFuture);
                        }
                        
                        return CompletableFuture.allOf(indexFutures.toArray(new CompletableFuture[0]));
                    } catch (Exception e) {
                        System.err.println("Failed to index directory " + dirPath + ": " + e.getMessage());
                        return CompletableFuture.completedFuture(null);
                    }
                });
    }

    /**
     * 查找符合条件的文件
     * 
     * @param dirPath 目录路径
     * @param includePatterns 包含的文件模式
     * @param excludePatterns 排除的文件模式
     * @return 文件路径列表
     */
    private List<String> findFiles(String dirPath, List<String> includePatterns, List<String> excludePatterns) {
        List<String> filePaths = new ArrayList<>();
        findFilesRecursive(new File(dirPath), filePaths, includePatterns, excludePatterns);
        return filePaths;
    }

    /**
     * 递归查找符合条件的文件
     * 
     * @param dir 目录
     * @param filePaths 文件路径列表
     * @param includePatterns 包含的文件模式
     * @param excludePatterns 排除的文件模式
     */
    private void findFilesRecursive(File dir, List<String> filePaths, List<String> includePatterns, List<String> excludePatterns) {
        if (!dir.isDirectory()) {
            return;
        }
        
        File[] files = dir.listFiles();
        if (files == null) {
            return;
        }
        
        for (File file : files) {
            if (file.isDirectory()) {
                // 检查目录是否应该被排除
                if (shouldExclude(file.getName(), excludePatterns)) {
                    continue;
                }
                
                findFilesRecursive(file, filePaths, includePatterns, excludePatterns);
            } else {
                String filePath = file.getAbsolutePath();
                
                // 检查文件是否应该被包含
                if (shouldInclude(filePath, includePatterns) && !shouldExclude(filePath, excludePatterns)) {
                    filePaths.add(filePath);
                }
            }
        }
    }

    /**
     * 检查文件是否应该被包含
     * 
     * @param filePath 文件路径
     * @param includePatterns 包含的文件模式
     * @return 是否应该被包含
     */
    private boolean shouldInclude(String filePath, List<String> includePatterns) {
        if (includePatterns == null || includePatterns.isEmpty()) {
            return true;
        }
        
        for (String pattern : includePatterns) {
            if (filePath.matches(pattern) || filePath.endsWith(pattern)) {
                return true;
            }
        }
        
        return false;
    }

    /**
     * 检查文件是否应该被排除
     * 
     * @param filePath 文件路径
     * @param excludePatterns 排除的文件模式
     * @return 是否应该被排除
     */
    private boolean shouldExclude(String filePath, List<String> excludePatterns) {
        if (excludePatterns == null || excludePatterns.isEmpty()) {
            return false;
        }
        
        for (String pattern : excludePatterns) {
            if (filePath.matches(pattern) || filePath.contains(pattern)) {
                return true;
            }
        }
        
        return false;
    }

    /**
     * 搜索代码
     * 
     * @param query 查询文本
     * @param limit 结果数量限制
     * @return 搜索结果列表的 Future
     */
    public CompletableFuture<List<vectordb.VectorSearchResult>> search(String query, int limit) {
        return this.init()
                .thenCompose(result -> this.embedding.embed(query))
                .thenCompose(embeddingVector -> this.vectorDB.search(this.collectionName, embeddingVector.getVector(), limit));
    }

    /**
     * 搜索代码（带过滤条件）
     * 
     * @param query 查询文本
     * @param limit 结果数量限制
     * @param filter 过滤条件
     * @return 搜索结果列表的 Future
     */
    public CompletableFuture<List<vectordb.VectorSearchResult>> search(String query, int limit, String filter) {
        return this.init()
                .thenCompose(result -> this.embedding.embed(query))
                .thenCompose(embeddingVector -> this.vectorDB.search(this.collectionName, embeddingVector.getVector(), limit, filter));
    }

    /**
     * 获取文档数量
     * 
     * @return 文档数量的 Future
     */
    public CompletableFuture<Long> count() {
        return this.init()
                .thenCompose(result -> this.vectorDB.count(this.collectionName));
    }

    /**
     * 关闭索引器
     * 
     * @return 完成后的 Future
     */
    public CompletableFuture<Void> close() {
        return this.vectorDB.close();
    }
}